!! Copyright (C) 2009,2010,2011,2012  Marco Restelli
!!
!! This file is part of:
!!   LDGH -- Local Hybridizable Discontinuous Galerkin toolkit
!!
!! LDGH is free software: you can redistribute it and/or modify it
!! under the terms of the GNU General Public License as published by
!! the Free Software Foundation, either version 3 of the License, or
!! (at your option) any later version.
!!
!! LDGH is distributed in the hope that it will be useful, but WITHOUT
!! ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
!! or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public
!! License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with LDGH. If not, see <http://www.gnu.org/licenses/>.
!!
!! author: Marco Restelli                   <marco.restelli@gmail.com>


!>\brief
!! This module collects some useful subroutines for handling the left
!! hand side of a linear system, mostly in combination with iterative
!! solvers.
!!
!! \n
!!
!! \todo This module must be completely reworked. Basically, it should
!! provide two functionalities:
!! <ul>
!!  <li> a general interface to the linear solver, so that when a
!!  matrix is provided it is then used to define an iterative solver,
!!  a MUMPS solver and so on. The idea should be simplifying the user
!!  interface at the cost of some more computations and additional
!!  working variables (the user will be able to avoid this by
!!  accessing directly mod_iterativesolvers_base, mod_mumpsintf and so
!!  on)
!!  <li> some simple preconditioners for iterative solvers.
!! </ul>
!!
!!
!! The utilities provided here fall into the following categories:
!! <ul>
!!  <li> <em>initialization</em>: \c mod_lhs_constructor, \c
!!  mod_lhs_destructor, \c lhs_setup;
!!  <li> <em>lhs evaluation (for iterative solvers)</em>: \c lhs;
!!  <li> <em>preconditioning</em>: \c precon_diag, \c precon_tridiag,
!!  \c precon_multidiag.
!! </ul>
!<----------------------------------------------------------------------
module mod_lhs

!-----------------------------------------------------------------------

 use mod_messages, only: &
   mod_messages_initialized, &
   error,   &
   warning, &
   info

 !$ use mod_omp_utils, only: &
 !$   mod_omp_utils_initialized, &
 !$   detailed_timing_omp, &
 !$   omput_push_key,      &
 !$   omput_pop_key,       &
 !$   omput_start_timer,   &
 !$   omput_close_timer,   &
 !$   omput_write_time

 use mod_kinds, only: &
   mod_kinds_initialized, &
   wp

 use mod_sparse, only: &
   mod_sparse_initialized, &
   ! sparse types
   t_col,       &
   transpose,   &
   matmul,      &
   diag,        &
   spdiag,      &
   clear

 use mod_umfintf, only: &
   mod_umfintf_initialized, &
   umfpack_factor, &
   umfpack_solve,  &
   umfpack_clean

!-----------------------------------------------------------------------
 
 implicit none

!-----------------------------------------------------------------------

! Module interface

 public :: &
   mod_lhs_constructor, &
   mod_lhs_destructor,  &
   mod_lhs_initialized, &
   lhs_setup,      &
   lhs_clean,      &
   lhs,            &
   precon_diag,    &
   precon_tridiag, &
   precon_multidiag

 private

!-----------------------------------------------------------------------

! Module types and parameters

! Module variables

 ! public members
 logical, protected ::               &
   mod_lhs_initialized = .false.
 ! private members
 type(t_col), save :: mt ! transposed matrix
 ! d0 is a two dimensional array for convenience in the function call
 real(wp), allocatable :: d0(:,:), d3(:,:), aa(:), bb(:), cc(:)
 type(t_col), save :: dn
 logical :: umf_need_clean

 character(len=*), parameter :: &
   this_mod_name = 'mod_lhs'

!-----------------------------------------------------------------------

contains

!-----------------------------------------------------------------------

 !> Constructor.
 !!
 !! \note If \c preconditioner is "multidiagonal", mod_umfintf must be
 !! initialized before calling this constructor, otherwise such module
 !! is ignored.
 !<
 subroutine mod_lhs_constructor(preconditioner)
  character(len=*), intent(in) :: preconditioner

  character(len=*), parameter :: &
    this_sub_name = 'constructor'

   !Consistency checks ---------------------------
   if( (mod_messages_initialized.eqv..false.) .or. &
!$     ( (detailed_timing_omp.eqv..true.).and. &
!$       (mod_omp_utils_initialized.eqv..false.) ) .or. &
       (mod_kinds_initialized.eqv..false.)    .or. &
       (mod_sparse_initialized.eqv..false.) ) then
     call error(this_sub_name,this_mod_name, &
                'Not all the required modules are initialized.')
   endif
   if( (mod_umfintf_initialized.eqv..false.) .and. &
       (trim(preconditioner).eq.'multidiagonal') )then
     call error(this_sub_name,this_mod_name, &
                'The required module "mod_umfintf" is not initialized.')
   endif
   if(mod_lhs_initialized.eqv..true.) then
     call warning(this_sub_name,this_mod_name, &
                  'Module is already initialized.')
   endif
   !----------------------------------------------

   mod_lhs_initialized = .true.
 end subroutine mod_lhs_constructor

!-----------------------------------------------------------------------
 
 !> Module destructor.
 !<
 subroutine mod_lhs_destructor()
  character(len=*), parameter :: &
    this_sub_name = 'destructor'
   
   !Consistency checks ---------------------------
   if(mod_lhs_initialized.eqv..false.) then
     call error(this_sub_name,this_mod_name, &
                'This module is not initialized.')
   endif
   !----------------------------------------------

   mod_lhs_initialized = .false.
 end subroutine mod_lhs_destructor

!-----------------------------------------------------------------------

 !> This subroutine sets two groups of module variables used in the
 !! iterative solution of a linear system.
 !! <ul>
 !!  <li> The sparse matrix \c mt holds the transposed system matrix:
 !!  due to the column oriented storage of sparse matrix, in the
 !!  iterative solution of the linear system the form \f$x^T\,A^T\f$
 !!  is much more efficient than the algebraically equivalent
 !!  \f$A\,x\f$.
 !!  <li> All the module matrices connected with the preconditioner.
 !! </ul>
 !! This function should be called right before starting the iterative
 !! solution of the linear system; if a given system has to be solved
 !! for many different right hand sides this function should be called
 !! only once before starting the solution of the systems, since it
 !! only depends on the system matrix and not on the right hand side.
 !!
 !! The optional parameter \c transposed_mat, when <tt>.true.</tt>,
 !! indicates that the matrix \c m is already provided as the
 !! transpose of the system matrix. This is useful since building
 !! directly the transposed matrix can be much more efficient that
 !! building the matrix and then transposing it.
 !!
 !! \warning Remember to call \c lhs_clean once you are done with a
 !! linear system and before calling this subroutine for a new one.
 !!
 !! \bug Several improvements are possible in this subroutine:
 !! <ul>
 !!   <li> if the system matrix is already available as \f$A^T\f$,
 !!   avoid making a copy in this module;
 !!   <li> Include some reordering utility to reduce the bandwidth;
 !!   <li> add other preconditioners, such as ILU.
 !! </ul>
 !<
 subroutine lhs_setup(m,preconditioner,ndiags,transposed_mat)
  !> The system matrix \f$A\f$.
  type(t_col), intent(in) :: m
  !> The preconditioner, currently implemented are \em diagonal, \em
  !! tridiagonal and \em multidiagonal. Any other value is simply
  !! ignored.
  !<
  character(len=*), intent(in) :: preconditioner
  !> Used only if <tt>preconditioner.eq.'multidiagonal'</tt>:
  !! specifies how many diagonal are used in the preconditioner. If
  !! <tt>ndiags.lt.0</tt> or if <tt>ndiags.gt.m\%n</tt> the whole
  !! matrix is used. This can be useful when using umfpack in double
  !! precision as preconditioner for an iterative solver working in
  !! quadruple precision.
  !<
  integer, intent(in), optional :: ndiags
  !> Pass this argument and set it to <tt>.true.</tt> to indicate that
  !! \c m is already transopsed.o
  !<
  logical, intent(in), optional :: transposed_mat
 
  logical :: need_transp
  integer :: i
  character(len=*), parameter :: &
    this_sub_name = 'lhs_setup'

   !Consistency checks ---------------------------
   if(mod_lhs_initialized.eqv..false.) then
     call error(this_sub_name,this_mod_name, &
                'This module is not initialized.')
   endif
   !----------------------------------------------

   need_transp = .true.
   if(present(transposed_mat)) then
     if(transposed_mat) need_transp = .false.
   endif
   if(need_transp) then
     !$ if(detailed_timing_omp) then
     !$   call omput_push_key("transpose")
     !$   call omput_start_timer()
     !$ endif
     mt = transpose(m)
     !$ if(detailed_timing_omp) then
     !$   call omput_write_time()
     !$   call omput_close_timer()
     !$   call omput_pop_key()
     !$ endif
   else
     mt = m
   endif
   umf_need_clean = .false.

   !$ if(detailed_timing_omp) then
   !$   call omput_push_key("preconditioner")
   !$   call omput_start_timer()
   !$ endif
   prec: select case(trim(preconditioner))

    case('diagonal')
     if(allocated(d0)) deallocate(d0)
     allocate(d0(m%n,1))
     d0 = 1.0_wp/diag(m,(/ 0 /))

    case('tridiagonal')
     if(allocated(d3)) deallocate(d3,aa,bb,cc)
     allocate(d3(m%n,3) , aa(m%n) , bb(m%n) , cc(m%n))
     d3 = diag(m,(/ -1 , 0 , 1 /))
     aa(1) = d3(1,2)
     cc(1) = d3(1,3)
     do i=2,m%n
       cc(i) = d3(i,3)
       bb(i) = d3(i-1,1)/aa(i-1)
       aa(i) = d3(i,2) - bb(i)*cc(i-1)
     enddo

    case('multidiagonal')
     ! get the diagonals 
     if(ndiags.ge.0) then
       dn = spdiag(m,(/ (i , i=-ndiags,ndiags) /))
     else
       dn = m
     endif
     call umfpack_factor(dn)
     umf_need_clean = .true.

   end select prec
   !$ if(detailed_timing_omp) then
   !$   call omput_write_time()
   !$   call omput_close_timer()
   !$   call omput_pop_key()
   !$ endif
 
 end subroutine lhs_setup
 
!-----------------------------------------------------------------------

 !> This subroutine should be called once for each call to \c
 !! lhs_setup to clean the allocated variables. It should be called
 !! when the system build with the call to \c lhs_setup is no longer
 !! needed.
 !<
 subroutine lhs_clean()

   call clear(mt)

   if(allocated(d0)) deallocate(d0)
   if(allocated(d3)) deallocate(d3,aa,bb,cc)

   if(umf_need_clean) call umfpack_clean()
   call clear(dn)

 end subroutine lhs_clean
 
!-----------------------------------------------------------------------

 !> This function evaluates the lhs as \f$lhs = A\,x\f$.
 pure &
 !$ subroutine omp_dummy_1(); end subroutine omp_dummy_1
 function lhs(x) result(mx)
  real(wp), intent(in) :: x(:)
  real(wp) :: mx(size(x))
 
   mx = matmul(x,mt)
 end function lhs
 
!-----------------------------------------------------------------------

 !> Precondition with the matrix diagonal (does not require umfpack).
 pure subroutine precon_diag(r)
  real(wp), intent(inout) :: r(:)
 
  r = d0(:,1)*r
 end subroutine precon_diag
 
!-----------------------------------------------------------------------
 
 !> Precondition with the three main diagonals (does not require
 !! umfpack).
 !<
 pure subroutine precon_tridiag(r)
  real(wp), intent(inout) :: r(:)

  integer :: i
  real(wp) :: y(size(r))
 
   y(1) = r(1)
   do i=2,size(r)
     y(i) = r(i) - bb(i)*y(i-1)
   enddo

   r(size(r)) = y(size(r))/aa(size(r))
   do i=size(r)-1,1,-1
     r(i) = (y(i) - cc(i)*r(i+1))/aa(i)
   enddo
 end subroutine precon_tridiag
 
!-----------------------------------------------------------------------
 
 !> Precondition with \c r main diagonals (requires umfpack).
 subroutine precon_multidiag(r)
  real(wp), intent(inout) :: r(:)

  real(wp) :: r_new(size(r))

  call umfpack_solve(dn,r_new,r)
  r = r_new

 end subroutine precon_multidiag
 
!-----------------------------------------------------------------------

end module mod_lhs

